package raft

import (
	"encoding/json"
	"fmt"
	"sync"
)

// KVStateMachine 实现简单的键值存储状态机
type KVStateMachine struct {
	mu    sync.Mutex
	store map[string]string
}

func NewKVStateMachine() *KVStateMachine {
	return &KVStateMachine{
		store: make(map[string]string),
	}
}

func (kv *KVStateMachine) Apply(cmd []byte) ([]byte, error) {
	var command struct {
		Op    string `json:"op"`
		Key   string `json:"key"`
		Value string `json:"value,omitempty"`
	}

	if err := json.Unmarshal(cmd, &command); err != nil {
		return nil, err
	}

	kv.mu.Lock()
	defer kv.mu.Unlock()

	switch command.Op {
	case "set":
		kv.store[command.Key] = command.Value
		return []byte("OK"), nil
	case "get":
		val, ok := kv.store[command.Key]
		if !ok {
			return nil, fmt.Errorf("key not found")
		}
		return []byte(val), nil
	case "del":
		delete(kv.store, command.Key)
		return []byte("OK"), nil
	default:
		return nil, fmt.Errorf("unknown operation")
	}
}

func (kv *KVStateMachine) Snapshot() ([]byte, error) {
	kv.mu.Lock()
	defer kv.mu.Unlock()
	return json.Marshal(kv.store)
}

func (kv *KVStateMachine) ApplySnapshot(snapshot []byte) error {
	kv.mu.Lock()
	defer kv.mu.Unlock()
	return json.Unmarshal(snapshot, &kv.store)
}
